#include <iostream>
#include <winsock2.h>
#include <windows.h>
#include <time.h>
#include <string>
#include <sstream>
#include <iomanip>

#pragma comment(lib, "ws2_32.lib")

#define NTP_PORT 123
#define NTP_PACKET_SIZE 48
#define NTP_TIMESTAMP_DELTA 2208988800ull

bool parseDateTime(const std::string &datetime, time_t &result)
{
    struct tm tm = {0};
    std::istringstream iss(datetime);
    char discard;

    iss >> tm.tm_year >> discard >> tm.tm_mon >> discard >> tm.tm_mday >> tm.tm_hour >> discard >> tm.tm_min >> discard >> tm.tm_sec;

    if (iss.fail())
        return false;

    tm.tm_year -= 1900;
    tm.tm_mon -= 1;
    tm.tm_isdst = -1;

    result = mktime(&tm);
    return result != -1;
}

void fillNtpPacket(char *packet, time_t customTime)
{
    memset(packet, 0, NTP_PACKET_SIZE);
    packet[0] = 0x24;

    time_t now = customTime + (time(NULL) - time(NULL));
    uint32_t ntpTime = htonl((uint32_t)(now + NTP_TIMESTAMP_DELTA));

    memcpy(&packet[40], &ntpTime, sizeof(ntpTime));
    memcpy(&packet[32], &ntpTime, sizeof(ntpTime));
    memcpy(&packet[16], &ntpTime, sizeof(ntpTime));
}

int main(int argc, char *argv[])
{
    if (argc != 2)
    {
        std::cerr << "Usage: " << argv[0] << " \"YYYY-MM-DD HH:MM:SS\"" << std::endl;
        return 1;
    }

    time_t customTime;
    if (!parseDateTime(argv[1], customTime))
    {
        std::cerr << "Invalid date format. Use \"YYYY-MM-DD HH:MM:SS\"" << std::endl;
        return 1;
    }

    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        std::cerr << "WSAStartup failed" << std::endl;
        return 1;
    }

    SOCKET sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET)
    {
        std::cerr << "Socket creation failed" << std::endl;
        WSACleanup();
        return 1;
    }

    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(NTP_PORT);
    serverAddr.sin_addr.s_addr = htonl(INADDR_ANY);

    if (bind(sock, (sockaddr *)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR)
    {
        std::cerr << "Bind failed" << std::endl;
        closesocket(sock);
        WSACleanup();
        return 1;
    }

    std::cout << "NTP server started with custom time: " << argv[1] << std::endl;

    char buffer[NTP_PACKET_SIZE];
    sockaddr_in clientAddr;
    int clientAddrLen = sizeof(clientAddr);

    while (true)
    {
        int bytesReceived = recvfrom(sock, buffer, NTP_PACKET_SIZE, 0, (sockaddr *)&clientAddr, &clientAddrLen);
        if (bytesReceived == SOCKET_ERROR)
        {
            std::cerr << "recvfrom failed" << std::endl;
            break;
        }

        char clientIP[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &clientAddr.sin_addr, clientIP, INET_ADDRSTRLEN);
        std::cout << "NTP request from " << clientIP << std::endl;

        fillNtpPacket(buffer, customTime);
        sendto(sock, buffer, NTP_PACKET_SIZE, 0, (sockaddr *)&clientAddr, clientAddrLen);
    }

    closesocket(sock);
    WSACleanup();
    return 0;
}